{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.Load Task in Benchmark Quickly"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Tutorial one is to load a task/env conveniently by simply giving task name and desired target embodiment (default is Franka)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "from PIL import Image\n",
    "from dm_control import viewer\n",
    "from VLABench.envs import load_env\n",
    "from VLABench.robots import *\n",
    "from VLABench.tasks import *\n",
    "from VLABench.configs import name2config\n",
    "from VLABench.utils.utils import find_key_by_value"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set the embodiment and task "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "task = \"select_fruit\"\n",
    "robot = \"franka\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If run this tutorial on server instead of PC, run in headless mode."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"MUJOCO_GL\"] = \"egl\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "###  1.1Load task and build environment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each time you load the single env, each instance varies in a large range."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = load_env(task, robot=robot, time_limit=1000)\n",
    "env.reset()\n",
    "\n",
    "image = env.render(camera_id=2, width=640, height=640)\n",
    "Image.fromarray(image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "obs=env.get_observation()\n",
    "obs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.2Run in interactive viewer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "viewer.launch(env)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.Overview of Primitive Tasks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import textwrap\n",
    "\n",
    "def plot_images(images, instructions, titles, max_columns=5, max_text_width=30):\n",
    "    assert len(images) == len(instructions) == len(titles), \"images should have the same length as instructions and titles\"\n",
    "    num_images = len(images)\n",
    "    num_rows = (num_images + max_columns - 1) // max_columns\n",
    "    fig, axes = plt.subplots(num_rows, max_columns, figsize=(max_columns * 3, num_rows * 3))\n",
    "    axes = axes.flatten()\n",
    "    for i, (image, instruction, title) in enumerate(zip(images, instructions, titles)):\n",
    "        if np.max(image) <= 1:\n",
    "            image = (image * 255).astype(np.uint8)\n",
    "        ax = axes[i]\n",
    "        ax.imshow(image)\n",
    "        ax.axis('off')\n",
    "        ax.set_title(title, fontsize=14)\n",
    "        wrapped_instruction = textwrap.fill(instruction, width=max_text_width)\n",
    "        ax.text(0.5, -0.1, wrapped_instruction, ha='center', va='top', transform=ax.transAxes, fontsize=12)\n",
    "        \n",
    "    for i in range(num_images, len(axes)):\n",
    "        axes[i].axis('off')\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_env_and_render(task, robot, **kwargs):\n",
    "    env = load_env(task, robot=robot, **kwargs)\n",
    "    env.reset()\n",
    "\n",
    "    image = env.render(camera_id=2, width=480, height=480)\n",
    "    instruction = env.task.get_instruction()\n",
    "    return image, instruction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1 Basic Scene: Mesh&Texture "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Default Setting: no additional domain randomization, no texture and scene augmenation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tasks_to_load = [\"select_toy\", \"select_fruit\", \"select_chemistry_tube\", \"select_mahjong\", \"select_poker\", \"add_condiment\", \"insert_flower\", \"select_book\", \"select_billiards\", \"select_drink\", \"select_ingredient\", \"select_painting\",\"put_box_on_painting\"]\n",
    "images, instructions = [], []\n",
    "for task in tasks_to_load:\n",
    "    try:\n",
    "        image, instruction = load_env_and_render(task, robot, reset_wait_step=0)\n",
    "        images.append(image)\n",
    "        instructions.append(instruction)\n",
    "    except Exception as e:\n",
    "        import traceback\n",
    "        traceback.print_exc()\n",
    "        print(f\"Failed to load task {task}: {e}\")\n",
    "plot_images(images, instructions, titles=tasks_to_load, max_columns=3)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2 Common Sense & World Knowledge"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The base environment for Common Sense & World Knowledge type tasks quering LLM to generate task instructions by default. We also recommend to generate the instructions in batch later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tasks_to_load = [\"hammer_loose_nail\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.3 Spatial Understanding"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The base environment for spatial understanding generate the relative spatial relationships between target entity and other entities by task-specific rules."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tasks_to_load = [\"select_toy_spatial\", \"select_fruit_spatial\", \"select_chemistry_tube_spatial\", \"select_mahjong_spatial\", \"select_poker_spatial\", \"add_condiment_spatial\", \"insert_flower_spatial\", \"select_book_spatial\", \"select_billiards_spatial\", \"select_drink_spatial\", \"select_ingredient_spatial\"]\n",
    "# \"\", \"select_poker\", \"add_condiment\", \"insert_flower\", \"select_book\", \"select_billiards\", \"select_drink\", \"select_ingredient\", \"select_painting\",\"put_box_on_painting\"\n",
    "images, instructions = [], []\n",
    "robot=\"franka\"\n",
    "for task in tasks_to_load:\n",
    "    try:\n",
    "        image, instruction = load_env_and_render(task, robot, reset_wait_step=0)\n",
    "        images.append(image)\n",
    "        instructions.append(instruction)\n",
    "    except Exception as e:\n",
    "        import traceback\n",
    "        traceback.print_exc()\n",
    "        print(f\"Failed to load task {task}: {e}\")\n",
    "plot_images(images, instructions, titles=tasks_to_load, max_columns=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.4 Semantic Understanding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tasks_to_load = [\"select_toy_semantic\", \"select_fruit_semantic\"]\n",
    "for task in tasks_to_load:\n",
    "    try:\n",
    "        image, instruction = load_env_and_render(task, robot, reset_wait_step=0)\n",
    "        images.append(image)\n",
    "        instructions.append(instruction)\n",
    "    except Exception as e:\n",
    "        import traceback\n",
    "        traceback.print_exc()\n",
    "        print(f\"Failed to load task {task}: {e}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "instructions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3.Overview of Composite Tasks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.1Cluster series tasks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tasks_to_load = [\"cluster_book\", \"cluster_billiards\", \"cluster_toy\", \"cluster_dessert\", \"cluster_drink\", \"cluster_ingredients\"]\n",
    "images, instructions = [], []\n",
    "for task in tasks_to_load:\n",
    "    image, instruction = load_env_and_render(task, robot, reset_wait_step=0)\n",
    "    images.append(image)\n",
    "    instructions.append(instruction)\n",
    "plot_images(images, instructions, titles=tasks_to_load, max_columns=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.2 Other Composite Tasks "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tasks_to_load = [\"texas_holdem\", \"texas_holdem_explore\", \"set_dining_table\", \"set_dining_left_hand\", \"set_dining_chopstick\", \"find_unseen_object\", \"cool_drink\", \"take_out_cool_drink\", \"book_rearrange\", \"heat_food\", \"rearrange_tube\", \"take_chemistry_experiment\", \"get_coffee\", \"get_coffee_with_sugar\", \"get_coffee_with_milk\", \"hammer_nail_and_hang_picture\", \"make_juice\", \"cook_dishes\", \"store_food\", \"play_mahjong\", \"complex_seesaw_use\", \"play_math_game\", \"set_study_table\"]\n",
    "images, instructions = [], []\n",
    "for task in tasks_to_load:\n",
    "    image, instruction = load_env_and_render(task, robot, reset_wait_step=0)\n",
    "    images.append(image)\n",
    "    instructions.append(instruction)\n",
    "plot_images(images, instructions, titles=tasks_to_load, max_columns=3)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "vlabench",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
